"""
In this file, we TODO
"""

import paddle
import paddle.nn as nn
from paddle.vision.models import resnet152

import pdb

class EncoderCNN(nn.Layer):
    def __init__(self, embed_size=256):
        """
        Load the pretrained ResNet-152 and replace the top FC layer for transfer learning.

        embed_size: int Dimension of word embedding vectors
        """
        super(EncoderCNN, self).__init__()
        resnet = resnet152(pretrained=True)
        modules = list(resnet.children())[:-1]
        self.resnet = nn.Sequential(*modules)
        # self.linear = nn.Linear(resnet.fc.in_features, embed_size)
        self.linear = nn.Linear(512*1, embed_size)  # A magic number found in paddlepaddle source code.
        self.bn = nn.BatchNorm1D(embed_size, momentum=0.01)

    def forward(self, images):
        """
        Extract feature vectors from input images.
        """
        # This raises an exception.
        pdb.set_trace()
        with paddle.no_grad():
            features = self.resnet(images)
            features = features.reshape(features.size(0), -1)
            features = self.linear(features)
            features = self.bn(features)
        return features

class DecoderRNN(nn.Layer):
    def __init__(self, embed_size, hidden_size, vocab_size, num_layers, max_seq_length=20):
        """Set the hyper-paremeters and build the layers"""
        super(DecoderRNN, self).__init__();
        self.embed = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers)
        self.linear = nn.Linear(hidden_size, vocab_size)
        self.max_seq_length = max_seq_length

    def forward(self, features, captions, lengths):
        """Decode image feature vectors and generates captions."""
        embeddings = self.embed(captions)
        embeddings = paddle.concat((paddle.unsqueeze(features,1), embeddings), 1)
        hiddens, _ = self.lstm(embeddings, sequence_length=lengths)
        outputs = self.linear(hiddens[0])
        return outputs

    def sample(self, features, states=None):
        """Generage captions for given image features using greedy search."""
        sampled_ids = []
        inputs = paddle.unsqueeze(features, 1)
        for i in range(self.max_seq_length):
            hiddens, states = self.lstm(inputs, states)
            outputs = self.linear(paddle.unsqueeze(hiddens, 1))
            _, predicted = paddle.max(outputs, 1)
            sampled_ids.append(predicted)
            inputs = self.embed(predicted)
            inputs = paddle.unsqueeze(inputs, 1)
        sampled_ids = paddle.stack(sampled_ids, 1)
        return sampled_ids

